#include <iostream>

#include <lemon/smart_graph.h>
#include <lemon/network_simplex.h>

#include "mincost_ns.hpp"

using namespace lemon;


using LGraph = SmartDigraph;

using Node = LGraph::Node;
using Arc = LGraph::Arc;
using ArcName = edge_name_t;

template<typename ValueType>
using ArcMap = SmartDigraph::ArcMap<ValueType>;

using Name = vertex_name_t;
template<typename ValueType>
using NodeMap = SmartDigraph::NodeMap<ValueType>;

using NS = NetworkSimplex<SmartDigraph, Capacity, Weight>;

Node add_node(LGraph& g, NodeMap<Name>& nmap, Name n) {
	Node u = g.addNode();
	nmap[u] = n;
	return u;
}

Arc add_arc(LGraph&g, ArcMap<Capacity>& capacities, ArcMap<Capacity>& lowerBounds, ArcMap<Weight>& weights, ArcMap<Capacity>& flows, ArcMap<ArcName>& arcs
		, Node o, Node d, Capacity c, Capacity lb, Weight w, Capacity f, ArcName id) {
	Arc a = g.addArc(o, d);
	capacities[a] = c;
	lowerBounds[a] = lb;
	weights[a] = w;
	flows[a] = f;
	arcs[a] = id;
	return a;
}

int mincost_ns(std::vector<edge_t>&edges, const nametable<vertex_name_t>& vertex_nt, vertex_name_t s, vertex_name_t t, Capacity requiredFlow, Weight& totalCost) {
    LGraph g;

    //add nodes
    std::vector<Node> nodes;
    nametable<Node> node_nt;

    for (auto& n: vertex_nt.as_vec()) {
    	Node u = g.addNode();
    	nodes.emplace_back(u);
    }
    node_nt.from_vec(nodes);

    Node s_ = node_nt[vertex_nt(s)];
    Node t_ = node_nt[vertex_nt(t)];

    //add arcs
    ArcMap<ArcName> arcs(g);
    ArcMap<Capacity> capacities(g);
    ArcMap<Capacity> lowerBounds(g);
    ArcMap<Weight> weights(g);
    ArcMap<Capacity> flows(g);  //may not be able use existing flows as input in lemon

    std::map<Arc, int> edgeMap;
    //for (auto& e: edges) {
    for (int i = 0; i < edges.size(); i++) {
    	auto& e = edges[i];
    	Arc a = add_arc(g, capacities, lowerBounds, weights, flows, arcs
    			, node_nt[vertex_nt(std::get<0>(e))], node_nt[vertex_nt(std::get<1>(e))]
    			, std::get<2>(e), std::get<3>(e), std::get<4>(e), std::get<5>(e), std::get<6>(e));
    	edgeMap[a] = i;
    }

    NS ns(g);
    ns.costMap(weights).upperMap(capacities).lowerMap(lowerBounds).stSupply(s_, t_, requiredFlow);

//    NS ns(g, capacities, weights, s, t, requiredFlow);

    NS::ProblemType status = ns.run();

    if (status == NS::OPTIMAL) {
    	totalCost = ns.totalCost();

    	//copy the result flows:
    	ns.flowMap(flows);
    	for (const auto& [a,i] : edgeMap) {
    		std::get<5>(edges[i]) = ns.flow(a);
    	}

    	return 0;
    }
    else if (status == NS::INFEASIBLE)
    	return -1;
    else //if (status == NS::UNBOUNDED)
    	return 1;
}


void test1() {
    LGraph g;
    Node u = g.addNode();
    Node v = g.addNode();
    Arc a = g.addArc(u, v);
    Arc b = g.addArc(u, v);

    ArcMap<Weight> weights(g);
    weights[a] = 5;
    weights[b] = 2;

    ArcMap<Capacity> capacities(g);
    capacities[a] = 10;
    capacities[b] = 5;

    NS ns(g);
    ns.costMap(weights).upperMap(capacities).stSupply(u, v, 12);

    ArcMap<Capacity> flows(g);
    NS::ProblemType status = ns.run();
    switch (status) {
    case NS::INFEASIBLE:
        std::cerr << "insufficient flow" << std::endl;
        break;
    case NS::OPTIMAL:
        ns.flowMap(flows);
        std::cerr << "flow[a]=" << ns.flow(a) << " flow[b]=" << flows[b] << std::endl;
        std::cerr << "cost=" << ns.totalCost() << std::endl; 
        break;
    case NS::UNBOUNDED:
        std::cerr << "infinite flow" << std::endl;
        break;
    default:
        break;
    }

    return;
}

/*
int main() {
  test1();
  
  return 0;
}
//*/
